MLX backend POC #1365

williambdean · 2025-04-11T15:39:43Z

Description

Getting ball rolling started with #1350

Related Issue

Closes #
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pytensor--1365.org.readthedocs.build/en/1365/

pytensor/link/pytorch/linker.py

codecov · 2025-04-11T18:11:07Z

Codecov Report

Attention: Patch coverage is 80.23256% with 17 lines in your changes missing coverage. Please review.

Project coverage is 82.04%. Comparing base (4e59f21) to head (d057453).
Report is 13 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/mlx/linker.py	77.50%	6 Missing and 3 partials ⚠️
pytensor/link/mlx/dispatch/basic.py	80.64%	6 Missing ⚠️
pytensor/compile/mode.py	50.00%	1 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (80.23%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1365      +/-   ##
==========================================
+ Coverage   82.02%   82.04%   +0.01%     
==========================================
  Files         203      208       +5     
  Lines       48845    48949     +104     
  Branches     8691     8701      +10     
==========================================
+ Hits        40067    40162      +95     
- Misses       6627     6632       +5     
- Partials     2151     2155       +4

Files with missing lines	Coverage Δ
pytensor/link/mlx/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/mlx/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/mlx/dispatch/math.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/linker.py	`100.00% <ø> (ø)`
pytensor/compile/mode.py	`84.09% <50.00%> (-0.64%)`	⬇️
pytensor/link/mlx/dispatch/basic.py	`80.64% <80.64%> (ø)`
pytensor/link/mlx/linker.py	`77.50% <77.50%> (ø)`

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

ricardoV94 · 2025-04-12T12:39:06Z

I suggest basing yourself on the numba linker, torch has a lot of hacks we hopefully don't need here

williambdean · 2025-04-12T20:29:04Z

Thanks for the pointer. I simplified the one method. Do you think that gen_functors can be removed as well? The only commonality with pytorch then is that no input can be numpy array.

ricardoV94 · 2025-04-13T13:19:54Z

Yeah you shouldn't need that you just need a call to tipify on the runtime inputs as well

pytensor/compile/mode.py

pytensor/link/mlx/linker.py

…to pr/1365

A shout out for the fathers of the day! Co-Authored-By: Ricardo Vieira <[email protected]> Co-Authored-By: Jesse Grabowski <[email protected]>

williambdean · 2025-04-19T15:42:46Z

Still need to get this to run:

import pytensor

pytensor.config.mode = "MLX"

Working

cetagostini · 2025-04-23T08:16:45Z

Hey big thanks to @jessegrabowski and @ricardoV94 to help with this PR!

I feel the PR is huge enough. Should we make a first merge and start to iterate on next versions? Cleaning and making all more consistent with other backends.

Thanks to @williambdean to open the PR!

ricardoV94 · 2025-04-24T09:46:39Z

pytensor/link/mlx/dispatch/basic.py

+
+
+@mlx_typify.register(np.ndarray)
+@mlx_typify.register(mx.array)


mxarray should be registered in mlx_typify_no_conversion_needed

ricardoV94 · 2025-04-24T09:47:06Z

pytensor/link/mlx/dispatch/basic.py

+@mlx_funcify.register(Assert)
+@mlx_funcify.register(CheckAndRaise)
+def mlx_funcify_CheckAndRaise(op, **kwargs):
+    warnings.warn(
+        f"""Skipping `CheckAndRaise` Op (assertion: {op.msg}) as MLX tracing would remove it.""",
+        stacklevel=2,
+    )
+
+    def assert_fn(x, *inputs):
+        return x
+
+    return assert_fn


Is this true, or just copy/pasta from JAX?

I need to check more here!

ricardoV94 · 2025-04-24T09:47:35Z

pytensor/link/mlx/dispatch/blockwise.py

+from pytensor.tensor.signal.conv import Conv1d
+
+
+def blockwise_conv1d(op, node, **kwargs):


Not needed anymore since they fixed upstream right?

I think we needed still. Where do you see its fixed? We are using this blockwise conv1d.

Sure but blockwise will call vmap on the core op, so we only need to dispatch core Conv1D to MLX Conv1D, then the blockwise variant will work automatically

ricardoV94 · 2025-04-24T09:48:08Z

pytensor/link/mlx/dispatch/core.py

+# ------------------------------------------------------------------
+# Join
+# ------------------------------------------------------------------


Don't introduce comments like this, the code is readable and won't get stale if we move things around

ricardoV94 · 2025-04-24T09:48:25Z

pytensor/link/mlx/dispatch/core.py

+# ------------------------------------------------------------------
+# Join
+# ------------------------------------------------------------------
+@mlx_funcify.register(Join)  # MLX


The multiple #MLX comments are useless

ricardoV94 · 2025-04-24T09:50:45Z

pytensor/link/mlx/dispatch/elemwise.py

+        # Convert scalar to array if needed
+        if isinstance(x, int | float) or (
+            isinstance(x, np.number) and not isinstance(x, np.ndarray)
+        ):
+            x = mx.array(x)


Should not be needed

ricardoV94 · 2025-04-24T09:51:27Z

pytensor/link/mlx/dispatch/elemwise.py

+@mlx_funcify.register(CAReduce)
+def mlx_funcify_CAReduce(op, **kwargs):
+    if isinstance(op.scalar_op, Add):
+
+        def sum(x):
+            return mx.sum(x, axis=op.axis)
+
+        return sum
+    elif isinstance(op.scalar_op, Mul):
+
+        def prod(x):
+            return mx.prod(x, axis=op.axis)
+
+        return prod
+    elif isinstance(op.scalar_op, AND):
+
+        def all(x):
+            return x.all(axis=op.axis)
+
+        return all
+    elif isinstance(op.scalar_op, OR):
+
+        def any(x):
+            return mx.any(x, axis=op.axis)
+
+        return any
+    else:
+        raise NotImplementedError(f"MLX does not support Elemwise {op.scalar_op}")


This should do a secon-lever dispatch on the core_op. Something like MLX_funcify_CAREduce that is called on op.scalar_op

ricardoV94 · 2025-04-24T09:53:14Z

pytensor/link/mlx/dispatch/math.py

+
+
+@mlx_funcify.register(Elemwise)
+def mlx_funcify_Elemwise(op, **kwargs):


Like CAReduce it should have a second level dispatch. Also we need to enforce the runtime_broadcastable checks (same in Alloc). And we shoud have a default implementation for that second level dispatch that tries to use getattr(MLX, "func_name") similar to how JAX does it already.

ricardoV94 · 2025-04-24T09:53:55Z

pytensor/link/mlx/dispatch/subtensor.py

+            if not op.inplace:
+                x = deepcopy(x)
+            x[indices] = y
+            return x


Need tests for all these including inplace variants

ricardoV94 · 2025-04-24T09:54:22Z

pytensor/link/mlx/linker.py

+        for n in self.fgraph.inputs:
+            sinput = storage_map[n]
+            # Handle random number generators specially
+            if isinstance(sinput[0], RandomState | Generator):
+                new_value = mlx_typify(
+                    sinput[0], dtype=getattr(sinput[0], "dtype", None)
+                )
+                sinput[0] = new_value
+            thunk_inputs.append(sinput)


Since we don't have Random stuff yet we shouldn't include the code

Okay added some!

ricardoV94 · 2025-06-03T08:27:35Z

pytensor/link/mlx/dispatch/blockwise.py

+    # 1) If it's a Conv1d Blockwise, use the custom implementation
+    if isinstance(op.core_op, Conv1d):
+        return blockwise_conv1d(op, node, **kwargs)


Here, we don't need this special casing anymore

ricardoV94 · 2025-06-03T08:27:50Z

pytensor/link/mlx/dispatch/core.py

+def mlx_funcify_Join(op, **kwargs):
+    def join(axis, *tensors):
+        view = op.view
+        if (view != -1) and all(


view was removed from PyTensor, ignore

ricardoV94 · 2025-06-03T08:30:18Z

pytensor/link/mlx/dispatch/math.py

+@mlx_typify.register(int)
+@mlx_typify.register(float)
+def mlx_typify_python_scalar(data, **kwargs):
+    return mx.array(data)


Put this close to the other mlx_typipy

ricardoV94 · 2025-06-03T08:31:29Z

pytensor/link/mlx/dispatch/math.py

+    def elemwise(*inputs):
+        # Enforce runtime broadcast checks (same as JAX and PyTorch implementations)
+        if node is not None:
+            # Convert inputs to MLX arrays for broadcast checking


node is never none, conversion to mx.array() shouldn't be needed either

ricardoV94 · 2025-06-03T08:32:59Z

tests/link/mlx/test_shape.py

+from tests.link.mlx.test_basic import compare_mlx_and_py
+
+
+@pytest.mark.xfail(reason="Shape Op is not supported yet")


Seems like it is?

ricardoV94 · 2025-06-03T08:34:09Z

pytensor/link/mlx/dispatch/elemwise.py

+    return min_reduce
+
+
+@mlx_funcify.register(CAReduce)


Put this before the specific scalar_op dispatches?

ricardoV94 · 2025-06-03T08:34:35Z

pytensor/link/mlx/dispatch/elemwise.py

+    return softmax_grad
+
+
+@mlx_funcify.register(Softplus)


Delete this? You have one in elemwise already

williambdean · 2025-06-03T11:51:57Z

Were we going to split this PR up into core functionality and op implementations? What are the next steps here?

ricardoV94 · 2025-06-03T11:59:19Z

This PR seems okay state that we can merge as is when the comments are addressed

williambdean added 2 commits April 11, 2025 11:38

mlx poc

d25f214

add test for dot

edacc0e

williambdean marked this pull request as draft April 11, 2025 15:43

williambdean commented Apr 11, 2025

View reviewed changes

pytensor/link/pytorch/linker.py Outdated Show resolved Hide resolved

williambdean added 10 commits April 11, 2025 12:26

restore pytorch

052fdc2

wrap in mx.array

a9ecad0

modify the pytorch jit

e690bff

move file

ad29c17

dont wrap

ba29b37

attempt to fix github action

8716870

change the rtol

9bf7edf

add init file

96ba116

skip if not installed

e116fa1

remove torch related code / comments

5d5f754

williambdean marked this pull request as ready for review April 11, 2025 17:54

williambdean added 2 commits April 12, 2025 16:25

simplify the fgraph_convert

b8cee3f

assert type

d057453

ricardoV94 reviewed Apr 13, 2025

View reviewed changes

pytensor/compile/mode.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Apr 13, 2025

View reviewed changes

pytensor/link/mlx/linker.py Outdated Show resolved Hide resolved

williambdean and others added 7 commits April 18, 2025 13:51

simplify the internal

ae202e6

remove the language

f1941fe

Adding operations in pytensor

7c8eae7

add extension

67a74fb

make compare function

fb5eb52

rename function

516b595

correct the function name

67bb8da

cetagostini and others added 11 commits April 18, 2025 20:13

Merge branch 'mlx-poc' of https://github.com/williambdean/pytensor in…

9d3eca8

…to pr/1365

Guys, I'm getting sad. We need help yisus!!!!!

9f31ab1

WILLIAM YOU NEED TO GO ANOTHER MILE! GO ON MY MATEEEEEEE, GO PHILLIES!

37440ff

RETURN, WHAT A SHAME! Sad times are coming.

4e4923f

AI COULD BE COOL? OR WE ARE JUST FUCKING AROUND?

6b27dc4

AI RULES BABY MY MATE

e308f83

test conv1d case

3744a18

I'm going for pizzas, it was an incredible day!

b41cab0

Merge branch 'mlx-poc' of https://github.com/williambdean/pytensor in…

323fa9d

…to pr/1365

SUUUUUUUUU!!!!!! LIFE IS GOING WELL. MLX FOR MEDIA MIX MODELS BAY

9766975

A shout out for the fathers of the day! Co-Authored-By: Ricardo Vieira <[email protected]> Co-Authored-By: Jesse Grabowski <[email protected]>

pre-commit

5ffc5ef

cetagostini added 2 commits April 19, 2025 11:32

Almost working

597f84e

Last PR sampling working

fb8fd2f

Working

ricardoV94 requested changes Apr 24, 2025

View reviewed changes

ricardoV94 added enhancement New feature or request backend compatibility labels Apr 24, 2025

cetagostini mentioned this pull request Apr 28, 2025

Computational overhead when transforming to a numpy array ml-explore/mlx#2132

Closed

cetagostini added 2 commits June 2, 2025 20:45

Requested changes by Ricardo

6a2b774

Pre commit changes

602f0ed

ricardoV94 reviewed Jun 3, 2025

View reviewed changes



		@mlx_typify.register(np.ndarray)
		@mlx_typify.register(mx.array)

		from pytensor.tensor.signal.conv import Conv1d


		def blockwise_conv1d(op, node, **kwargs):



		@mlx_funcify.register(Elemwise)
		def mlx_funcify_Elemwise(op, **kwargs):

		from tests.link.mlx.test_basic import compare_mlx_and_py


		@pytest.mark.xfail(reason="Shape Op is not supported yet")

MLX backend POC #1365

Are you sure you want to change the base?

MLX backend POC #1365

Conversation

williambdean commented Apr 11, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Checklist

Type of change

Uh oh!

Uh oh!

codecov bot commented Apr 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 commented Apr 12, 2025

Uh oh!

williambdean commented Apr 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Apr 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

williambdean commented Apr 19, 2025

Uh oh!

cetagostini commented Apr 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

williambdean commented Apr 11, 2025 •

edited by github-actions bot

Loading

codecov bot commented Apr 11, 2025 •

edited

Loading

williambdean commented Apr 12, 2025 •

edited

Loading

ricardoV94 commented Apr 13, 2025 •

edited

Loading

ricardoV94 commented Jun 3, 2025 •

edited

Loading